#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
PostgreSQL维护模块
提供数据库维护功能：清理、重建索引等
"""

import os
import re
from ...core.utils import logger, run_command, get_input_with_default, confirm_action

class PostgresqlMaintenance:
    """PostgreSQL维护功能类"""
    
    def vacuum_database(self):
        """对数据库执行VACUUM FULL操作"""
        logger.info("运行VACUUM FULL以清理数据库")
        
        try:
            run_command(["sudo", "-u", "postgres", "psql", "-c", "VACUUM FULL;"])
            logger.info("VACUUM完成")
        except Exception as e:
            logger.error(f"执行VACUUM操作时出错: {e}")
    
    def reindex_database(self):
        """对数据库重建索引"""
        db_name = get_input_with_default("输入要重新索引的数据库名称", "")
        
        if not db_name:
            logger.error("数据库名称不能为空")
            return
            
        logger.info(f"重新索引数据库 {db_name}...")
        
        try:
            run_command(["sudo", "-u", "postgres", "psql", "-c", f"REINDEX DATABASE {db_name};"])
            logger.info("重新索引已完成")
        except Exception as e:
            logger.error(f"重建索引时出错: {e}")
    
    def update_postgresql_conf(self, parameter, value, comment=None):
        """
        更新postgresql.conf文件中的参数
        
        Args:
            parameter: 要更新的参数名
            value: 参数值
            comment: 可选的注释
        
        Returns:
            bool: 更新是否成功
        """
        if not self.config_path:
            logger.error("找不到PostgreSQL配置路径")
            return False
            
        postgresql_conf = os.path.join(self.config_path, "postgresql.conf")
        
        if not os.path.exists(postgresql_conf):
            logger.error(f"配置文件 {postgresql_conf} 不存在")
            return False
            
        try:
            # 读取文件内容
            with open(postgresql_conf, "r") as file:
                lines = file.readlines()
            
            # 检查参数是否已存在
            param_pattern = re.compile(rf'^[#\s]*{re.escape(parameter)}\s*=')
            param_exists = False
            param_uncommented = False
            
            for i, line in enumerate(lines):
                if param_pattern.match(line):
                    param_exists = True
                    if not line.strip().startswith('#'):
                        param_uncommented = True
                    break
            
            # 如果参数是shared_preload_libraries，需要特殊处理
            if parameter == "shared_preload_libraries":
                if param_exists:
                    # 检查是否已经包含pg_stat_statements
                    for i, line in enumerate(lines):
                        if param_pattern.match(line):
                            if "pg_stat_statements" in line:
                                logger.info("pg_stat_statements已在shared_preload_libraries中配置")
                                return True
                            
                            # 提取当前值并添加pg_stat_statements
                            current_value_match = re.search(r'=\s*[\'"](.+)[\'"]', line)
                            if current_value_match:
                                current_value = current_value_match.group(1)
                                libraries = [lib.strip() for lib in current_value.split(',')]
                                if "pg_stat_statements" not in libraries:
                                    libraries.append("pg_stat_statements")
                                    new_value = ", ".join(libraries)
                                    new_line = f"{parameter} = '{new_value}'"
                                    if comment:
                                        new_line += f"  # {comment}"
                                    new_line += "\n"
                                    lines[i] = new_line
                                    logger.info(f"已将pg_stat_statements添加到shared_preload_libraries中: {new_value}")
                            else:
                                # 如果无法提取当前值，替换为新值
                                new_line = f"{parameter} = '{value}'"
                                if comment:
                                    new_line += f"  # {comment}"
                                new_line += "\n"
                                lines[i] = new_line
                            break
                else:
                    # 如果参数不存在，添加新行
                    new_line = f"{parameter} = '{value}'"
                    if comment:
                        new_line += f"  # {comment}"
                    new_line += "\n"
                    lines.append("\n# 性能监控配置\n")
                    lines.append(new_line)
                    logger.info(f"已添加shared_preload_libraries = '{value}'配置")
            else:
                # 其他参数的常规处理
                if param_exists:
                    # 更新现有参数
                    for i, line in enumerate(lines):
                        if param_pattern.match(line):
                            new_line = f"{parameter} = {value}"
                            if comment:
                                new_line += f"  # {comment}"
                            new_line += "\n"
                            lines[i] = new_line
                            break
                else:
                    # 添加新参数
                    new_line = f"{parameter} = {value}"
                    if comment:
                        new_line += f"  # {comment}"
                    new_line += "\n"
                    lines.append(new_line)
            
            # 写回文件
            with open(postgresql_conf, "w") as file:
                file.writelines(lines)
                
            logger.info(f"已更新{postgresql_conf}文件")
            return True
                
        except Exception as e:
            logger.error(f"更新配置文件时出错: {e}")
            return False
    
    def enable_pg_stat_statements(self):
        """启用pg_stat_statements扩展"""
        logger.info("正在启用pg_stat_statements扩展...")
        
        # 1. 在postgresql.conf中添加shared_preload_libraries配置
        config_updated = self.update_postgresql_conf(
            "shared_preload_libraries", 
            "pg_stat_statements", 
            "启用性能监控扩展"
        )
        
        if config_updated:
            logger.info("已将pg_stat_statements添加到shared_preload_libraries配置")
            logger.warning("需要重启PostgreSQL服务才能生效")
            
            # 询问是否重启PostgreSQL服务
            if confirm_action("是否立即重启PostgreSQL服务?"):
                try:
                    self.restart()
                    logger.info("PostgreSQL服务已重启")
                except Exception as e:
                    logger.error(f"重启PostgreSQL服务时出错: {e}")
        
        # 2. 在数据库中创建扩展
        try:
            run_command(["sudo", "-u", "postgres", "psql", "-c", "CREATE EXTENSION IF NOT EXISTS pg_stat_statements;"])
            logger.info("pg_stat_statements扩展已成功启用")
        except Exception as e:
            logger.error(f"启用pg_stat_statements扩展时出错: {e}")
    
    def show_query_performance(self):
        """显示查询性能统计数据"""
        logger.info("正在显示查询性能统计数据...")
        
        try:
            run_command(["sudo", "-u", "postgres", "psql", "-c", "SELECT * FROM pg_stat_statements ORDER BY total_exec_time DESC LIMIT 20;"])
        except Exception as e:
            logger.error(f"显示查询性能统计数据时出错: {e}") 